import backtrader as bt

from 数据预处理模块 import 数据预处理类, 获取数据类


class 运行回测类:
    def __init__(self,回测策略,起始日期,结束日期,初始资金):
        self.回测策略=回测策略
        self.起始日期=起始日期
        self.结束日期=结束日期
        self.初始资金=初始资金


    def 运行(self):
        self.初始化cerebro()
        self.加载数据预处理()
        self.加载数据()
        self.添加策略()
        self.设置初始资金()
        self.打印当前资金('策略执行前资金：')
        self.运行策略()
        self.打印当前资金('策略执行后资金：')
        self.可视化()

    def 初始化cerebro(self):
        self.cerebro = bt.Cerebro()

    def 加载数据(self):
        self.cerebro.adddata(self.数据)
    def 添加策略(self):
        self.cerebro.addstrategy(self.回测策略)
    def 设置初始资金(self):
        self.cerebro.broker.setcash(self.初始资金)

    def 运行策略(self):
        return self.cerebro.run()
    def 可视化(self):
        self.cerebro.plot()

    def 打印当前资金(self,txt):
        print(txt+' %.2f' % self.cerebro.broker.getvalue())
    def 加载数据预处理(self):
        a=获取数据类()
        a.获取tushare数据预设置()
        self.数据 = a.获取股票当日开盘信息(股票代码='600515.SH', 起始日期=self.起始日期, 结束日期=self.结束日期)
        b=数据预处理类(self.数据)

        b.数据预处理()

        self.数据=b.数据转换(开始日期=self.起始日期,结束日期=self.结束日期)



        return  self.数据












